#include "pool.h"
#include <algorithm>
#include <random>
#include <set>
#include <stdexcept>
#include <unordered_map>

bool Pool::init(const Topo& topo, int ec_data, int ec_parity) {
    ec_data_ = ec_data;
    ec_parity_ = ec_parity;
    all_disks_.clear();
    int server_idx = 0;
    for (const auto& s : topo.servers()) {
        int disk_idx = 0;
        for (const auto& d : s.disks) {
            all_disks_.emplace_back(server_idx, disk_idx);
            ++disk_idx;
        }
        ++server_idx;
    }
    int total_disks = all_disks_.size();
    int pg_count = total_disks * 20;
    pgs_.clear();
    // 统计每个盘已分配PG数
    std::vector<std::vector<int>> disk_pg_cnt(topo.servers().size());
    for (size_t i = 0; i < topo.servers().size(); ++i)
        disk_pg_cnt[i].resize(topo.servers()[i].disks.size(), 0);
    std::mt19937 rng{std::random_device{}()};
    int server_num = topo.servers().size();
    int frag_num = ec_data_ + ec_parity_;
    for (int i = 0; i < pg_count; ++i) {
        PG pg;
        // 1. 随机打乱server顺序
        std::vector<int> server_indices(server_num);
        for (int s = 0; s < server_num; ++s) server_indices[s] = s;
        std::shuffle(server_indices.begin(), server_indices.end(), rng);
        int remain = frag_num;
        int k = 0;
        // 2. 先在不同server上各选一个盘
        for (; k < std::min(server_num, frag_num); ++k) {
            int sidx = server_indices[k];
            int min_pg = *std::min_element(disk_pg_cnt[sidx].begin(), disk_pg_cnt[sidx].end());
            std::vector<int> candidates;
            for (int d = 0; d < (int)disk_pg_cnt[sidx].size(); ++d) {
                if (disk_pg_cnt[sidx][d] == min_pg)
                    candidates.push_back(d);
            }
            std::uniform_int_distribution<> dist(0, candidates.size()-1);
            int pick_disk = candidates[dist(rng)];
            disk_pg_cnt[sidx][pick_disk]++;
            pg.disk_locs.push_back({sidx, pick_disk});
        }
        remain -= k;
        // 3. 如果还需要更多分片，允许重复server
        while (remain > 0) {
            // 在所有server中选PG数最少的盘
            int min_pg = INT32_MAX;
            for (size_t s = 0; s < disk_pg_cnt.size(); ++s)
                min_pg = std::min(min_pg, *std::min_element(disk_pg_cnt[s].begin(), disk_pg_cnt[s].end()));
            std::vector<std::pair<int, int>> candidates;
            for (size_t s = 0; s < disk_pg_cnt.size(); ++s) {
                for (size_t d = 0; d < disk_pg_cnt[s].size(); ++d) {
                    if (disk_pg_cnt[s][d] == min_pg)
                        candidates.emplace_back(s, d);
                }
            }
            std::uniform_int_distribution<> dist(0, candidates.size()-1);
            auto pick = candidates[dist(rng)];
            disk_pg_cnt[pick.first][pick.second]++;
            pg.disk_locs.push_back({(int)pick.first, (int)pick.second});
            --remain;
        }
        pgs_.push_back(pg);
    }
    pg_chunk_count_.assign(pg_count, 0);
    chunkid_meta_.clear();
    return true;
}

uint64_t Pool::alloc_chunk_id() {
    return chunk_id_++;
}

int Pool::allocated_chunk_count() const {
    return chunk_id_ - 1;
}

std::vector<uint64_t> Pool::allocated_chunks() const {
    std::vector<uint64_t> res;
    for (uint64_t i = 1; i < chunk_id_; ++i) res.push_back(i);
    return res;
}

ChunkMeta Pool::alloc_chunkid_balanced() {
    // 选择chunk数最少的PG
    int min_pg = *std::min_element(pg_chunk_count_.begin(), pg_chunk_count_.end());
    std::vector<int> candidates;
    for (int i = 0; i < (int)pg_chunk_count_.size(); ++i) {
        if (pg_chunk_count_[i] == min_pg) candidates.push_back(i);
    }
    std::mt19937 rng{std::random_device{}()};
    std::uniform_int_distribution<> dist(0, candidates.size()-1);
    int target_pg = candidates[dist(rng)];
    uint64_t cid = chunk_id_++;
    pg_chunk_count_[target_pg]++;
    ChunkMeta meta{cid, target_pg};
    chunkid_meta_[cid] = meta;
    return meta;
}

const ChunkMeta* Pool::get_chunk_meta(uint64_t chunkid) const {
    auto it = chunkid_meta_.find(chunkid);
    if (it != chunkid_meta_.end()) return &it->second;
    return nullptr;
}

const PG* Pool::get_pg(int pgid) const {
    if (pgid < 0 || pgid >= (int)pgs_.size()) return nullptr;
    return &pgs_[pgid];
}

int Pool::get_pg_of_chunk(uint64_t chunkid) const {
    auto it = chunkid_meta_.find(chunkid);
    if (it != chunkid_meta_.end()) return it->second.pg_id;
    return -1;
} 